{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "635d8ebb",
   "metadata": {},
   "source": [
    "# Ensemble Retriever with Convex Combination (CC)\n",
    "\n",
    "- Author: [Harheem Kim](https://github.com/harheem)\n",
    "- Peer Review:\n",
    "- Proofread : [JaeJun Shim](https://github.com/kkam-dragon)\n",
    "- This is a part of [LangChain Open Tutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/10-Retriever/11-Convex-Combination-Ensemble-Retriever.ipynb)[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/10-Retriever/11-Convex-Combination-Ensemble-Retriever.ipynb)\n",
    "## Overview\n",
    "\n",
    "This tutorial focuses on implementing and comparing different ensemble retrieval methods in LangChain. While LangChain's built-in ```EnsembleRetriever``` uses the **Reciprocal Rank Fusion (```RRF```)** method, we'll explore an additional approach by implementing the **Convex Combination (```CC```)** method.<br>\n",
    "The tutorial guides you through creating custom implementations of both **```RRF``` and ```CC``` methods** , allowing for a direct performance comparison between these ensemble techniques.\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "- [Overview](#overview)\n",
    "- [Environment Setup](#environment-setup)\n",
    "- [Process Document](#process-document)\n",
    "- [Initialize Retrievers](#initialize-retrievers)\n",
    "- [Implement Ensemble Retrievers](#implement-ensemble-retrievers)\n",
    "- [Compare and Test](#compare-and-test)\n",
    "\n",
    "### References\n",
    "\n",
    "- [EnsembleRetriever](https://python.langchain.com/api_reference/langchain/retrievers/langchain.retrievers.ensemble.EnsembleRetriever.html)\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6c7aba4",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials. \n",
    "- You can checkout the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21943adb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain-opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f25ec196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langchain_community\",\n",
    "        \"langchain_openai\",\n",
    "        \"langchain_core\",\n",
    "        \"faiss-cpu\",\n",
    "        \"pdfplumber\",\n",
    "        \"rank_bm25\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f9065ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment variables have been set successfully.\n"
     ]
    }
   ],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"OPENAI_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"Conversation-With-History\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "690a9ae0",
   "metadata": {},
   "source": [
    "You can alternatively set ```OPENAI_API_KEY``` in **.env** file and load it.\n",
    "\n",
    "[Note] This is not necessary if you've already set ```OPENAI_API_KEY``` in previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4f99b5b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load API key information\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "342b34aa",
   "metadata": {},
   "source": [
    "##  Process Document\n",
    "\n",
    "This section outlines the preparation process for processing PDF documents before storing them in a vector store. \n",
    "\n",
    "We use ```PDFPlumberLoader``` to load the PDF file and leverage ```RecursiveCharacterTextSplitter``` to break down the document into smaller, manageable chunks. \n",
    "\n",
    "The chunk size is set to 200 characters with no overlap, allowing for efficient processing while maintaining the document's semantic integrity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "564a7b34",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PDFPlumberLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "# Load the PDF document\n",
    "loader = PDFPlumberLoader(\"data/Introduction_LangChain.pdf\")\n",
    "# Split the document into manageable chunks\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=0)\n",
    "split_documents = loader.load_and_split(text_splitter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d90629",
   "metadata": {},
   "source": [
    "## Initialize Retrievers\n",
    "\n",
    "This section initializes retrievers to implement two different search approaches. We create embeddings using OpenAI's ```text-embedding-3-small``` model and set up ```FAISS``` vector search based on these embeddings. \n",
    "\n",
    "Additionally, we configure a ```BM25``` retriever for keyword-based search, with both retrievers set to return the top 5 most relevant results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "93539287",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_community.retrievers import BM25Retriever\n",
    "\n",
    "# Initialize OpenAI embeddings\n",
    "embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\n",
    "\n",
    "# Initialize FAISS retriever with vector embeddings\n",
    "faiss = FAISS.from_documents(\n",
    "    documents=split_documents, embedding=embeddings\n",
    ").as_retriever(search_kwargs={\"k\": 5})\n",
    "\n",
    "# Initialize BM25 retriever for keyword-based search\n",
    "bm25 = BM25Retriever.from_documents(documents=split_documents)\n",
    "bm25.k = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "809b08dc",
   "metadata": {},
   "source": [
    "## Implement Ensemble Retrievers\n",
    "\n",
    "This section introduces a custom retriever implementing two ensemble search methods, designed to compare performance against LangChain's built-in ```EnsembleRetriever``` . \n",
    "\n",
    "We implement both **Reciprocal Rank Fusion (```RRF```)** , which combines results based on document rankings, and **Convex Combination (```CC```)** , which utilizes normalized scores. \n",
    "\n",
    "Both methods integrate results from ```FAISS``` and ```BM25``` retrievers to provide more accurate and diverse search results, allowing users to select the most suitable ensemble approach for their needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dc5b9555",
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "from typing import List, Optional\n",
    "from langchain_core.retrievers import BaseRetriever\n",
    "from langchain_core.documents import Document\n",
    "from pydantic import BaseModel, model_validator\n",
    "\n",
    "\n",
    "class EnsembleMethod(str, Enum):\n",
    "    RRF = \"rrf\"  # Reciprocal Rank Fusion\n",
    "    CC = \"cc\"  # Convex Combination\n",
    "\n",
    "\n",
    "class EnsembleRetriever(BaseRetriever, BaseModel):\n",
    "    retrievers: List[BaseRetriever]\n",
    "    weights: Optional[List[float]] = None\n",
    "    method: EnsembleMethod = EnsembleMethod.RRF\n",
    "    c: int = 60\n",
    "\n",
    "    @model_validator(mode=\"before\")\n",
    "    def validate_weights(cls, values):\n",
    "        weights = values.get(\"weights\")\n",
    "        method = values.get(\"method\", EnsembleMethod.RRF)\n",
    "\n",
    "        if not weights:\n",
    "            n_retrievers = len(values[\"retrievers\"])\n",
    "            values[\"weights\"] = [1 / n_retrievers] * n_retrievers\n",
    "        elif method == EnsembleMethod.CC and abs(sum(weights) - 1.0) > 1e-6:\n",
    "            raise ValueError(\"CC method의 경우 weights의 합이 1이어야 합니다\")\n",
    "\n",
    "        return values\n",
    "\n",
    "    def _get_relevant_documents(self, query: str) -> List[Document]:\n",
    "        docs_list = [\n",
    "            retriever.get_relevant_documents(query) for retriever in self.retrievers\n",
    "        ]\n",
    "\n",
    "        if self.method == EnsembleMethod.RRF:\n",
    "            return self._rrf_fusion(docs_list)\n",
    "        else:\n",
    "            return self._cc_fusion(docs_list)\n",
    "\n",
    "    def _rrf_fusion(self, docs_list: List[List[Document]]) -> List[Document]:\n",
    "        \"\"\"\n",
    "        Implements Reciprocal Rank Fusion algorithm\n",
    "        - Combines results based on document rankings\n",
    "        - Uses a constant 'c' to prevent high ranks from dominating\n",
    "        - Applies weights to different retrievers' contributions\n",
    "        \"\"\"\n",
    "        from collections import defaultdict\n",
    "\n",
    "        scores = defaultdict(float)\n",
    "        for docs, weight in zip(docs_list, self.weights):\n",
    "            for rank, doc in enumerate(docs, 1):\n",
    "                scores[doc.page_content] += weight / (rank + self.c)\n",
    "\n",
    "        all_docs = []\n",
    "        seen = set()\n",
    "        for docs in docs_list:\n",
    "            for doc in docs:\n",
    "                if doc.page_content not in seen:\n",
    "                    all_docs.append(doc)\n",
    "                    seen.add(doc.page_content)\n",
    "\n",
    "        return sorted(all_docs, key=lambda x: scores[x.page_content], reverse=True)\n",
    "\n",
    "    def _cc_fusion(self, docs_list: List[List[Document]]) -> List[Document]:\n",
    "        \"\"\"\n",
    "        Implements Convex Combination fusion\n",
    "        - Combines normalized scores from different retrievers\n",
    "        - Requires weights to sum to 1.0\n",
    "        - Handles cases with missing or zero scores\n",
    "        \"\"\"\n",
    "        from collections import defaultdict\n",
    "\n",
    "        scores = defaultdict(float)\n",
    "        for docs, weight in zip(docs_list, self.weights):\n",
    "            max_score = max(\n",
    "                (doc.metadata.get(\"score\", 1.0) for doc in docs), default=1.0\n",
    "            )\n",
    "            if max_score == 0:\n",
    "                max_score = 1.0\n",
    "\n",
    "            for doc in docs:\n",
    "                norm_score = doc.metadata.get(\"score\", 1.0) / max_score\n",
    "                scores[doc.page_content] += weight * norm_score\n",
    "\n",
    "        all_docs = []\n",
    "        seen = set()\n",
    "        for docs in docs_list:\n",
    "            for doc in docs:\n",
    "                if doc.page_content not in seen:\n",
    "                    all_docs.append(doc)\n",
    "                    seen.add(doc.page_content)\n",
    "\n",
    "        return sorted(all_docs, key=lambda x: scores[x.page_content], reverse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ccbbc652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import EnsembleRetriever as OriginalEnsembleRetriever\n",
    "\n",
    "# Initialize the original LangChain EnsembleRetriever\n",
    "original_ensemble_retriever = OriginalEnsembleRetriever(retrievers=[faiss, bm25])\n",
    "\n",
    "# Initialize Ensemble Retriever with RRF (Reciprocal Rank Fusion) method\n",
    "rrf_ensemble_retriever = EnsembleRetriever(\n",
    "    retrievers=[faiss, bm25], method=EnsembleMethod.RRF\n",
    ")\n",
    "\n",
    "# Initialize Ensemble Retriever with CC (Convex Combination) method\n",
    "cc_ensemble_retriever = EnsembleRetriever(\n",
    "    retrievers=[faiss, bm25],\n",
    "    method=EnsembleMethod.CC,\n",
    "    weights=[0.5, 0.5],  # Equal weights for both retrievers\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "390c87bc",
   "metadata": {},
   "source": [
    "## Compare and Test\n",
    "\n",
    "This section presents a test function for comparing ensemble retrieval results. \n",
    "\n",
    "While the **```RRF``` method** , which follows LangChain's default implementation, produces identical results to **Original** , the **```CC``` method** utilizing normalized scores and weights offers different search patterns. \n",
    "\n",
    "By testing with real queries and comparing these approaches, we can identify which ensemble method better suits our project requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c33665c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pretty_print(query):\n",
    "    for i, (original_doc, cc_doc, rrf_doc) in enumerate(\n",
    "        zip(\n",
    "            original_ensemble_retriever.invoke(query),\n",
    "            cc_ensemble_retriever.invoke(query),\n",
    "            rrf_ensemble_retriever.invoke(query),\n",
    "        )\n",
    "    ):\n",
    "        print(f\"[{i}] [Original] Q: {query}\", end=\"\\n\\n\")\n",
    "        print(original_doc.page_content)\n",
    "        print(\"-\" * 100)\n",
    "        print(f\"[{i}] [RRF] Q: {query}\", end=\"\\n\\n\")\n",
    "        print(rrf_doc.page_content)\n",
    "        print(\"-\" * 100)\n",
    "        print(f\"[{i}] [CC] Q: {query}\", end=\"\\n\\n\")\n",
    "        print(cc_doc.page_content)\n",
    "        print(\"=\" * 100, end=\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4ee9ec4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "Introductions to all the key parts of LangChain you’ll need to know! Here you'll find high level\n",
      "explanations of all LangChain concepts.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[0] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "Introductions to all the key parts of LangChain you’ll need to know! Here you'll find high level\n",
      "explanations of all LangChain concepts.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[0] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "Introductions to all the key parts of LangChain you’ll need to know! Here you'll find high level\n",
      "explanations of all LangChain concepts.\n",
      "====================================================================================================\n",
      "\n",
      "[1] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "For a deeper dive into LangGraph concepts, check out this page.\n",
      "Integrations\n",
      "LangChain is part of a rich ecosystem of tools that integrate with our framework and build on top of it. If\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[1] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "For a deeper dive into LangGraph concepts, check out this page.\n",
      "Integrations\n",
      "LangChain is part of a rich ecosystem of tools that integrate with our framework and build on top of it. If\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[1] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "For a deeper dive into LangGraph concepts, check out this page.\n",
      "Integrations\n",
      "LangChain is part of a rich ecosystem of tools that integrate with our framework and build on top of it. If\n",
      "====================================================================================================\n",
      "\n",
      "[2] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "If you're looking to build something specific or are more of a hands-on learner, check out our tutorials\n",
      "section. This is the best place to get started.\n",
      "These are the best ones to get started with:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[2] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "If you're looking to build something specific or are more of a hands-on learner, check out our tutorials\n",
      "section. This is the best place to get started.\n",
      "These are the best ones to get started with:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[2] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "LangChain simplifies every stage of the LLM application lifecycle:\n",
      "Development: Build your applications using LangChain's open-source components and third-party\n",
      "====================================================================================================\n",
      "\n",
      "[3] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "Integration packages (e.g. , , etc.): Important\n",
      "langchain-openai langchain-anthropic\n",
      "integrations have been split into lightweight packages that are co-maintained by the LangChain\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[3] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "Integration packages (e.g. , , etc.): Important\n",
      "langchain-openai langchain-anthropic\n",
      "integrations have been split into lightweight packages that are co-maintained by the LangChain\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[3] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "The LangChain framework consists of multiple open-source libraries. Read more in the Architecture\n",
      "page.\n",
      ": Base abstractions for chat models and other components.\n",
      "langchain-core\n",
      "====================================================================================================\n",
      "\n",
      "[4] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "LangChain simplifies every stage of the LLM application lifecycle:\n",
      "Development: Build your applications using LangChain's open-source components and third-party\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[4] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "LangChain simplifies every stage of the LLM application lifecycle:\n",
      "Development: Build your applications using LangChain's open-source components and third-party\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[4] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "Read up on security best practices to make sure you're developing safely with LangChain.\n",
      "Contributing\n",
      "====================================================================================================\n",
      "\n",
      "[5] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      ": Third-party integrations that are community maintained.\n",
      "langchain-community\n",
      ": Orchestration framework for combining LangChain components into production-ready\n",
      "langgraph\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[5] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      ": Third-party integrations that are community maintained.\n",
      "langchain-community\n",
      ": Orchestration framework for combining LangChain components into production-ready\n",
      "langgraph\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[5] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "If you're looking to build something specific or are more of a hands-on learner, check out our tutorials\n",
      "section. This is the best place to get started.\n",
      "These are the best ones to get started with:\n",
      "====================================================================================================\n",
      "\n",
      "[6] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "The LangChain framework consists of multiple open-source libraries. Read more in the Architecture\n",
      "page.\n",
      ": Base abstractions for chat models and other components.\n",
      "langchain-core\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[6] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "The LangChain framework consists of multiple open-source libraries. Read more in the Architecture\n",
      "page.\n",
      ": Base abstractions for chat models and other components.\n",
      "langchain-core\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[6] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "Integration packages (e.g. , , etc.): Important\n",
      "langchain-openai langchain-anthropic\n",
      "integrations have been split into lightweight packages that are co-maintained by the LangChain\n",
      "====================================================================================================\n",
      "\n",
      "[7] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "Read up on security best practices to make sure you're developing safely with LangChain.\n",
      "Contributing\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[7] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "Read up on security best practices to make sure you're developing safely with LangChain.\n",
      "Contributing\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[7] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      ": Third-party integrations that are community maintained.\n",
      "langchain-community\n",
      ": Orchestration framework for combining LangChain components into production-ready\n",
      "langgraph\n",
      "====================================================================================================\n",
      "\n",
      "[8] [Original] Q: What are the advantages of LangChain?\n",
      "\n",
      "Head to the reference section for full documentation of all classes and methods in the LangChain\n",
      "Python packages.\n",
      "Ecosystem\n",
      "🦜🛠 LangSmith\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[8] [RRF] Q: What are the advantages of LangChain?\n",
      "\n",
      "Head to the reference section for full documentation of all classes and methods in the LangChain\n",
      "Python packages.\n",
      "Ecosystem\n",
      "🦜🛠 LangSmith\n",
      "----------------------------------------------------------------------------------------------------\n",
      "[8] [CC] Q: What are the advantages of LangChain?\n",
      "\n",
      "Head to the reference section for full documentation of all classes and methods in the LangChain\n",
      "Python packages.\n",
      "Ecosystem\n",
      "🦜🛠 LangSmith\n",
      "====================================================================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pretty_print(\"What are the advantages of LangChain?\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-opentutorial-ZA2wmMtu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
